#include <stdio.h>
#include <sys/types.h>
#include <sys/socket.h>
#include <netinet/in.h>
#include <arpa/inet.h>
#include <unistd.h>
#include <stdlib.h>
#include <string.h>
#include <sqlite3.h>
#include <pthread.h>  // 多线程处理升级

#define BUF_SIZE 4096  // 大缓冲区存升级文件
#define DB_NAME "UserInformation.db"
// 登录注册指令
#define REG_CMD "REG"         // 注册
#define LOGIN_CMD "LOGIN"     // 登录
#define REG_OK "REG_OK"       // 注册成功
#define REG_EXIST "REG_EXIST" // 账号已存在
#define LOGIN_OK "LOGIN_OK"   // 登录成功
#define LOGIN_ERR "LOGIN_ERR" // 登录失败
// 升级指令/端口
#define UPGRADE_CMD "UPGRADE"       // 升级指令（格式：UPGRADE|文件名）
#define UPGRADE_READY "UPGRADE_READY"// 准备接收文件
#define UPGRADE_SUCC "UPGRADE_SUCC"  // 升级成功
#define UPGRADE_FAIL "UPGRADE_FAIL"  // 升级失败
#define UPGRADE_PORT 8890            // 升级专用端口
#define MAIN_PORT 8889               // 主端口（登录/注册/设备地址）

// 全局变量
sqlite3 *sqldb = NULL;

// --------------------------
// 原有功能：数据库/登录/注册（精简保留）
// --------------------------
int db_init() 
{
    const char *create_sql = "CREATE TABLE IF NOT EXISTS UserInfoTable("
                             "id INTEGER PRIMARY KEY AUTOINCREMENT,"
                             "account TEXT NOT NULL UNIQUE,"
                             "password TEXT NOT NULL);";
    char *errmsg = NULL;
    
    if (sqlite3_open(DB_NAME, &sqldb) != SQLITE_OK) 
    {
        printf("数据库打开失败：%s\n", sqlite3_errmsg(sqldb));
        return -1;
    }
    
    if (sqlite3_exec(sqldb, create_sql, NULL, NULL, &errmsg) != SQLITE_OK) 
    {
        printf("建表失败：%s\n", errmsg);
        sqlite3_free(errmsg);
        return -1;
    }
    
    return 0;
}

void handle_reg(int clifd, char acc[], char pwd[]) 
{
    char sql[BUF_SIZE] = {0};
    char *errmsg = NULL;
    char **resultp = NULL;
    int nrow, ncol;

    snprintf(sql, BUF_SIZE, "SELECT account FROM UserInfoTable WHERE account='%s';", acc);
    
    if (sqlite3_get_table(sqldb, sql, &resultp, &nrow, &ncol, &errmsg) != SQLITE_OK) 
    {
        printf("注册查询失败：%s\n", errmsg);
        sqlite3_free(errmsg);
        return;
    }
    
    if (nrow > 0) 
    {
        send(clifd, REG_EXIST, strlen(REG_EXIST), 0);
        printf("账号已存在\n");
        return;
    }
    
    snprintf(sql, BUF_SIZE, "INSERT INTO UserInfoTable VALUES(NULL,'%s','%s');", acc, pwd);
    
    if (sqlite3_exec(sqldb, sql, NULL, NULL, &errmsg) != SQLITE_OK) 
    {
        printf("注册插入失败：%s\n", errmsg);
        sqlite3_free(errmsg);
        return;
    }
    
    send(clifd, REG_OK, strlen(REG_OK), 0);
    printf("注册成功\n");
}

void handle_login(int clifd, char acc[], char pwd[]) 
{
    char sql[BUF_SIZE] = {0};
    char *errmsg = NULL;
    char **resultp = NULL;
    int nrow, ncol;

    snprintf(sql, BUF_SIZE, "SELECT account FROM UserInfoTable WHERE account='%s' AND password='%s';", acc, pwd);
    
    if (sqlite3_get_table(sqldb, sql, &resultp, &nrow, &ncol, &errmsg) != SQLITE_OK) 
    {
        printf("登录查询失败：%s\n", errmsg);
        sqlite3_free(errmsg);
        send(clifd, LOGIN_ERR, strlen(LOGIN_ERR), 0);
        return;
    }
    
    if (nrow > 0) 
    {
        send(clifd, LOGIN_OK, strlen(LOGIN_OK), 0);
        printf("登录成功\n");
    } 
    else 
    {
        send(clifd, LOGIN_ERR, strlen(LOGIN_ERR), 0);
        printf("登录失败\n");
    }
}

// 新增：升级包接收+本地保存
void save_upgrade_file(int clifd, char *file_name) 
{
    if (!file_name) 
    {
        send(clifd, UPGRADE_FAIL, strlen(UPGRADE_FAIL), 0);
        printf("升级失败：文件名空\n");
        return;
    }

    // 回复“准备接收”
    send(clifd, UPGRADE_READY, strlen(UPGRADE_READY), 0);
    printf("准备接收升级包：%s\n", file_name);

    // 创建文件保存
    FILE *fp = fopen(file_name, "wb");
    
    if (!fp) 
    {
        perror("创建升级文件失败");
        send(clifd, UPGRADE_FAIL, strlen(UPGRADE_FAIL), 0);
        return;
    }

    // 分块接收文件
    char buf[BUF_SIZE] = {0};
    ssize_t recv_len = 0;
    
    while ((recv_len = recv(clifd, buf, BUF_SIZE, 0)) > 0) 
    {
        fwrite(buf, 1, recv_len, fp);
        memset(buf, 0, BUF_SIZE);
    }

    // 收尾
    fclose(fp);
    
    if (recv_len < 0) 
    {
        perror("接收升级包失败");
        remove(file_name); // 删除损坏文件
        send(clifd, UPGRADE_FAIL, strlen(UPGRADE_FAIL), 0);
    } 
    else 
    {
        send(clifd, UPGRADE_SUCC, strlen(UPGRADE_SUCC), 0);
        printf("升级包保存完成：%s\n", file_name);
    }
}

// 升级端口（8890）处理线程
void *upgrade_thread(void *arg) 
{
    (void)arg;
    int upgrade_fd = socket(AF_INET, SOCK_STREAM, 0);
    
    if (upgrade_fd < 0) 
    {
        perror("升级socket创建失败");
        pthread_exit(NULL);
    }

    // 绑定8890端口
    struct sockaddr_in serv_addr;
    serv_addr.sin_family = AF_INET;
    serv_addr.sin_port = htons(UPGRADE_PORT);
    serv_addr.sin_addr.s_addr = INADDR_ANY;
    int opt = 1;
    setsockopt(upgrade_fd, SOL_SOCKET, SO_REUSEADDR, &opt, sizeof(opt)); // 端口复用

    if (bind(upgrade_fd, (struct sockaddr *)&serv_addr, sizeof(serv_addr)) < 0) 
    {
        perror("升级bind失败");
        close(upgrade_fd);
        pthread_exit(NULL);
    }
    
    if (listen(upgrade_fd, 5) < 0) 
    {
        perror("升级listen失败");
        close(upgrade_fd);
        pthread_exit(NULL);
    }
    
    printf("升级服务启动：%d端口\n", UPGRADE_PORT);

    // 循环处理升级连接
    struct sockaddr_in cli_addr;
    socklen_t cli_len = sizeof(cli_addr);
    
    while (1) 
    {
        int clifd = accept(upgrade_fd, (struct sockaddr *)&cli_addr, &cli_len);
        
        if (clifd < 0) 
        {
            perror("升级accept失败");
            continue;
        }
        
        printf("新升级连接：%s:%d\n", inet_ntoa(cli_addr.sin_addr), ntohs(cli_addr.sin_port));

        // 解析升级指令（UPGRADE|文件名）
        char buf[BUF_SIZE] = {0};
        ssize_t len = recv(clifd, buf, BUF_SIZE-1, 0);
        
        if (len <= 0) 
        {
            close(clifd);
            continue;
        }
        
        char *cmd = strtok(buf, "|");
        char *file_name = strtok(NULL, "|");
        
        if (cmd && !strcmp(cmd, UPGRADE_CMD) && file_name) 
        {
            save_upgrade_file(clifd, file_name);
        } 
        else 
        {
            send(clifd, UPGRADE_FAIL, strlen(UPGRADE_FAIL), 0);
            printf("无效升级指令：%s\n", buf);
        }

        close(clifd);
        printf("升级连接断开\n");
    }
    
    close(upgrade_fd);
    pthread_exit(NULL);
}

// 主端口（8889）：登录/注册/设备地址
void handle_main_cmd(int clifd, char *buf) 
{
    char cmd[64] = {0}, acc[64] = {0}, pwd[64] = {0};
    char *p = strtok(buf, "|");
    
    if (!p) 
    {
        return;
    }
    strncpy(cmd, p, 63);

    // 处理登录/注册
    p = strtok(NULL, "|");
    
    if (!p) 
    {
        return;
    }
    
    strncpy(acc, p, 63);
    p = strtok(NULL, "|");
    
    if (!p) 
    {
        return;
    }
    
    strncpy(pwd, p, 63);

    if (!strcmp(cmd, REG_CMD)) 
    {
        handle_reg(clifd, acc, pwd);
    }
    else if (!strcmp(cmd, LOGIN_CMD)) 
    {
        handle_login(clifd, acc, pwd);
    }
}

int main() 
{
    // 初始化数据库
    if (db_init() != 0) 
    {
        printf("数据库初始化失败，退出\n");
        return -1;
    }

    // 启动升级线程
    pthread_t tid;
    
    if (pthread_create(&tid, NULL, upgrade_thread, NULL) != 0) 
    {
        perror("创建升级线程失败");
        sqlite3_close(sqldb);
        return -1;
    }
    
    pthread_detach(tid); // 分离线程，避免内存泄漏

    // 主端口（8889）UDP搜索
    int udp_fd = socket(AF_INET, SOCK_DGRAM, 0);
    
    if (udp_fd < 0) 
    {
        perror("UDP socket失败");
        sqlite3_close(sqldb);
        return -1;
    }
    
    struct sockaddr_in udp_addr;
    udp_addr.sin_family = AF_INET;
    udp_addr.sin_port = htons(MAIN_PORT);
    udp_addr.sin_addr.s_addr = INADDR_ANY;
    
    if (bind(udp_fd, (struct sockaddr *)&udp_addr, sizeof(udp_addr)) < 0) 
    {
        perror("UDP bind失败");
        close(udp_fd);
        sqlite3_close(sqldb);
        return -1;
    }

    // 接收UDP搜索（bilibili指令）
    char buf[BUF_SIZE] = {0};
    struct sockaddr_in cli_addr;
    socklen_t cli_len = sizeof(cli_addr);
    ssize_t len = recvfrom(udp_fd, buf, BUF_SIZE-1, 0, (struct sockaddr *)&cli_addr, &cli_len);
    
    if (len <= 0 || strcmp(buf, "bilibili") != 0) 
    {
        printf("无效UDP指令\n");
        close(udp_fd);
        sqlite3_close(sqldb);
        return -1;
    }
    
    sendto(udp_fd, "yes", 3, 0, (struct sockaddr *)&cli_addr, cli_len);
    printf("收到搜索指令，回复确认\n");
    close(udp_fd);

    // 主端口（8889）TCP处理登录/注册
    int tcp_fd = socket(AF_INET, SOCK_STREAM, 0);
    
    if (tcp_fd < 0) 
    {
        perror("TCP socket失败");
        sqlite3_close(sqldb);
        return -1;
    }
    
    struct sockaddr_in tcp_addr;
    tcp_addr.sin_family = AF_INET;
    tcp_addr.sin_port = htons(MAIN_PORT);
    tcp_addr.sin_addr.s_addr = INADDR_ANY;
    int opt=1;
    setsockopt(tcp_fd, SOL_SOCKET, SO_REUSEADDR, &opt, sizeof(opt));
    
    if (bind(tcp_fd, (struct sockaddr *)&tcp_addr, sizeof(tcp_addr)) < 0) 
    {
        perror("TCP bind失败");
        close(tcp_fd);
        sqlite3_close(sqldb);
        return -1;
    }
    
    if (listen(tcp_fd, 5) < 0) 
    {
        perror("TCP listen失败");
        close(tcp_fd);
        sqlite3_close(sqldb);
        return -1;
    }
    
    printf("主服务启动：%d端口(登录/注册）\n", MAIN_PORT);

    // 处理主连接
    int clifd = accept(tcp_fd, (struct sockaddr *)&cli_addr, &cli_len);
    
    if (clifd < 0) 
    {
        perror("TCP accept失败");
        close(tcp_fd);
        sqlite3_close(sqldb);
        return -1;
    }
    
    printf("新主连接：%s:%d\n", inet_ntoa(cli_addr.sin_addr), ntohs(cli_addr.sin_port));

    while (1) 
    {
        memset(buf, 0, BUF_SIZE);
        len = recv(clifd, buf, BUF_SIZE-1, 0);
        
        if (len <= 0) 
        {
            if (len < 0) 
            {
                perror("主连接recv失败");
            }
            else 
            {
                printf("主连接断开\n");
            }
            break;
        }
        
        handle_main_cmd(clifd, buf);
    }

    // 资源释放
    close(clifd);
    close(tcp_fd);
    sqlite3_close(sqldb);
    return 0;
}